# !/usr/bin/python3
# Created by Ross on 19-3-26
import os
from collections import Counter

import numpy as np
import pandas as pd
import tensorflow.contrib.eager as tfe
from keras.preprocessing.sequence import pad_sequences
from sklearn.manifold import TSNE

from hparams import Hparams
from model.JointModel import JointModel
from utils import load_hparams, get_checkpoints, id2label
import matplotlib
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from utils import id2label

hparams = Hparams()
hp = hparams.parser.parse_args()  # 超参数字典
# load_hparams(hp, hp.log_dir)
load_hparams(hp, 'log/2019-03-28-12:11')

test_x = np.load(os.path.join(hp.data_dir, 'test_x.npy'))
test_y = np.load(os.path.join(hp.data_dir, 'test_y.npy')).astype(np.int32)

test_lens = [min(len(x), hp.seq_maxlen) for x in test_x]
test_x = pad_sequences(test_x, hp.seq_maxlen, 'float32', padding='post', truncating='post')
if hp.fake_task == 'POS':
    test_POS = np.load(os.path.join(hp.data_dir, 'POStest_x.npy'))
    test_POS = pad_sequences(test_POS, hp.seq_maxlen, 'float32', padding='post', truncating='post')

with open('data/10fold_shijie_segment/test_x.txt', 'r', encoding='utf-8') as fp:
    sentences = [line.strip() for line in fp]

model = JointModel(hp.seq_maxlen, hp.emb_size, hp.rnn_size, hp.keep_prob,
                   hp.fake_task, hp.num_class, hp.use_crf, ntags=hp.ntags)
ckpts = get_checkpoints(hp.log_dir)
acc = tfe.metrics.Accuracy()
pred = []

#     for ckpt in ckpts:
#         acc.init_variables()
#         model.load_weights(ckpt)
#         domain_pred = model.predict_domain([test_x, test_lens])
#         pred.append(domain_pred.numpy())
#     result_and_model = []
#     for i in zip(*pred):
#         c = Counter(i)
#         result_and_model.append(c.most_common(1)[0][0])

pos, domain, h = model([test_x, test_lens], is_training=False, output_H=True)

# 降维到3-D
tsne = TSNE(n_components=3)
data = tsne.fit_transform(h)


def plot_embedding(data, id_, title):
    x_min, x_max = np.min(data, 0), np.max(data, 0)
    data = (data - x_min) / (x_max - x_min)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    for i in range(data.shape[0]):
        # if id_[i] == 3:
        #     continue
        ax.text(data[i, 0], data[i, 1], data[i, 2], id2label[id_[i]],
                color=plt.cm.Set2(id_[i] / 31.))
    plt.title(title)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_zlim(0, 1)
    return fig


print('plot')
fig = plot_embedding(data, test_y, 'output_h')
print('ready to show')
plt.show()
fig.savefig('output_h_visualization.png')
